import numcodecs
import zarr
from pathlib import Path
import numpy as np
import yaml
from pynput import keyboard
import time
from ezrobot.logger.data_logger import TrajectoryDataLogger

from ezrobot.utils import GraphRunner, HandEyeConfig
from ezrobot.sensor import ViveTracker, SensorManager, AngleEncoder, RealSense, TactileSensor
from ezrobot.utils.general import accurate_sleep
from ezgl import tb, Matrix4x4

ee2gripper = np.array([[1,0,0,0],
               [0,1,0,0],
               [0,0,1,0.11772],
               [0,0,0,1]])


BASE_DIR = Path(__file__).resolve().parent
compressor = numcodecs.Blosc(cname='zstd', clevel=5, shuffle=numcodecs.Blosc.SHUFFLE)
handeye_config = HandEyeConfig.load(BASE_DIR / "../configs/handeye_config.yaml")

def record_data(data, target_pose, config):
    global stored_data
    stored_data.setdefault('color', []).append(
        data[f"RealSense_{config['robot_settings']['realsense_ser_num']}"]['color'][:250,:].copy()
    )
    stored_data.setdefault('depth', []).append(
        data[f"RealSense_{config['robot_settings']['realsense_ser_num']}"]['depth'][:250,:].copy()
    )
    stored_data.setdefault('gripper_distance', []).append(
        data['AngleEncoder_1']
    )
    stored_data.setdefault('target_poses', []).append(
        target_pose.toTransRotV
        ector()
    )
    stored_data.setdefault('tactile_umi_left_src', []).append(
        data[f"TactileSensor_{config['UMI']['tactile_ser_num']['left']}"]['src']
    )
    stored_data.setdefault('tactile_umi_left_force', []).append(
        data[f"TactileSensor_{config['UMI']['tactile_ser_num']['left']}"]['force']
    )

def write_data(task_name: str = "p&p_cup_UMI"):
    global save_key
    global stored_data
    global traj_num
    if save_key:
         
        save_key = False
        file = zarr.open(rf"./data/a_data2/{task_name}_{traj_num}.zarr", mode="w")
        for key, value_list in stored_data.items():
            array = np.stack(value_list)
            file.create_dataset(key, data=array, compressor=compressor, chunks=True)
        print(f"data saved to {rf'./data/{task_name}_{traj_num}.zarr'}")
        
        # reset
        input("press enter to collect next!")
        save_key = False
        stored_data = dict()
        traj_num += 1
        print("buffer cleared")
        return True
    return False

def on_press(key):
    global save_key
    try:
        if str(key) == 'Key.space':
            save_key = True
    except Exception as e:
        raise e

def get_target_pose(data):
    """
    获取目标位姿, vive2hand
    """
    vive_pose = data['ViveTracker_1']
    return handeye_config.base2map * vive_pose * handeye_config.hand2eye.inverse() * Matrix4x4(ee2gripper)

def run(sensors, config, fps=50):
    t0 = time.monotonic()
    data = sensors.get_data()
    target_pose = get_target_pose(data)
    record_data(data, target_pose, config)

    elapsed_time = time.monotonic() - t0
    if 1/fps > elapsed_time:
        accurate_sleep(1/fps - elapsed_time)
    
    write_data()

if __name__ == '__main__':
    data_logger = TrajectoryDataLogger(path=r"C:\Users\LENOVO\Desktop\ezrobot\data")
    data_logger.start_trajectory()
    global save_key
    global stored_data

    global traj_num
    save_key = False
    traj_num = 0

    config_path = r"./configs/config.yaml"
    with open(config_path, 'r') as f:
            config = yaml.safe_load(f)

    # multi sensors
    encoder = AngleEncoder(config['UMI']["encoder_id"], config['UMI']["encoder_index"])
    vive = ViveTracker(1)
    tactile_UMI_l = TactileSensor(config['UMI']['tactile_ser_num']['left'])
    rs = RealSense(index=config['robot_settings']['realsense_ser_num'])
    sensors = SensorManager(sensor=[encoder, vive, tactile_UMI_l, rs], frequency=30, buffer_size=10, streaming=True)

    listener = keyboard.Listener(on_press=on_press)
    listener.start()

    input("Press Enter to start...")
    stored_data = dict()

    while 1:
        run(sensors, config)

